#    Copyright (C) 2003-2008 William Waites
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""
Code for retrieving information from (R)WHOIS servers.
"""
import logging
import socket
import bsddb
import re

__all__ = ["WhoisServer", "AutonomousSystemCache"]

class WhoisServer(object):
	"""
	Instances of this class, initialized with a hostname and an optional
	port provide a query method to retrieve information about a resource
	from the whois server. IPv4 and IPv6 connections are both fully
	supported.
	"""
	def __init__(self, host, port=43):
		self.host = host
		self.port = port
		self.log = logging.getLogger(self.host)
	def _socket(self):
		for result in socket.getaddrinfo(self.host, self.port,
						 socket.AF_UNSPEC, socket.SOCK_STREAM):
			af, type, proto, cname, addr = result
			self.log.debug("trying %s" % (addr[0],))
			sock = socket.socket(af, type, proto)
			try:
				sock.connect(addr)
				return sock
			except socket.gaierror:
				sock.close()
		raise socket.gaierror((8, 'could not connect to remote host'))
	def query(self, resource):
		self.log.debug("query: %s" % (resource,))
		try:
			sock = self._socket()
			sock.send("%s\r\n" % (resource,))
			fp = sock.makefile()
			result = fp.read()
			fp.close()
			return result
		except:
			self.log.error("error performing query", exc_info=True)
			raise

## TODO BUG w/ 18268 -- next level referral and ambiguous RPSL
class AutonomousSystemCache(object):
	"""
	Using the whois servers of the various RIRs, try to query information
	about a particular Autonomous System and cache the result. The constructor
	takes an optional argument for a filename in which to cache the results
	(defaults to "ascache.db" in the current directory").

	The interface is dictionary-like, for example,

	>>> cache = AutonomousSystemCache(fname=None)
	>>> print cache[239]
	{'name': 'UTORONTO-AS'}
	>>>
	"""
	comline = re.compile(r"^[#%].*$")
	kvpline = re.compile(r"^(?P<key>[A-Za-z\-]+):[^A-Za-z0-9]+(?P<value>.*)$")
	valcont = re.compile(r"^[ \t]+(?P<value>.*)$")
	asrange = re.compile(r"^[0-9]+ - [0-9]+$")
	def __init__(self, fname="ascache.db", debug=False):
		self.cache = bsddb.hashopen(fname, "c")
		self.servers = {
			"arin" : WhoisServer("whois.arin.net"),
			"ripe" : WhoisServer("whois.ripe.net"),
			"afnic" : WhoisServer("whois.afnic.net"),
			"apnic" : WhoisServer("whois.apnic.net"),
			"lacnic" : WhoisServer("whois.lacnic.net"),
		}
		if debug:
			self.log = logging.getLogger("whois")
		else:
			class Log(object):
				def debug(*av, **kw): pass
			self.log = Log()
	def parse(self, data):
		objects = []
		obj = None
		for line in map(lambda x: x.rstrip(), data.split("\n")):
			self.log.debug(line)
			if not line:
				obj = None
				continue
			m = self.comline.match(line)
			if m:
				continue
			m = self.kvpline.match(line)
			if m:
				if not obj:
					obj = {}
					objects.append(obj)
				gd = m.groupdict()
				key = gd["key"]
				value = gd["value"]
				obj[key] = value
				continue
			m = self.valcont.match(line)
			if m:
				if not obj: raise "ParseError"
				gd = m.groupdict()
				value = gd["value"]
				obj[key] = "\n".join((obj[key], value))
		return objects
	def __getitem__(self, key, flush=False):
		if not isinstance(key, str): key = str(key)
		if flush or key not in self.cache:
			data = self.servers["arin"].query(key)
			rec = self.parse(data)
			for obj in rec:
				if "OrgID" in obj:
					org = obj["OrgID"].lower()
					if org in self.servers:
						data = self.servers[org].query("AS%s" % (key,))
						rec = self.parse(data)
						break
			self.cache[key] = data
		else:
			data = self.cache[key]
			rec = self.parse(data)

		result = {'name' : 'AS%s' % (key,)}
		for obj in rec:
			if "aut-num" in obj:
				result.update(self.ripe_result(rec))
			elif "ASNumber" in obj:
				result.update(self.arin_result(rec))
		return result				

	def ripe_result(self, rec):
		result = {}
		asn = None
		for obj in rec:
			if "aut-num" in obj:
				asn = obj
				break
		if asn:
			for tag in ("as-name", "owner-c"):
				if tag in asn:
					result["name"] = asn[tag]
					break
		return result

	def arin_result(self, rec):
		result = {}
		asn = None
		for obj in rec:
			if "ASNumber" in obj:
				asn = obj
				break
		if asn:
			result["name"] = asn["ASName"]
		return result

if __name__ == '__main__':
	import doctest
	doctest.testmod()
